import numpy as np
import pandas as pd
from qtorch.core.strategy import Strategy

class ThresholdStrategy(Strategy):
    """阈值策略（基于预测值生成交易信号）"""
    def __init__(self, upper_threshold=0.015, lower_threshold=-0.015):
        self.upper = upper_threshold
        self.lower = lower_threshold
        
    def generate_signals(self, data: pd.DataFrame) -> pd.Series:
        required_columns = {'prediction', 'volatility'}
        missing = required_columns - set(data.columns)
        if missing:
            raise ValueError(f"Missing required columns: {', '.join(missing)}")
            
        # 处理缺失值
        data = data.fillna(0)
        
        # 动态调整阈值（结合波动率）
        volatility = data['volatility'].rolling(20).std().fillna(0)
        upper = self.upper * (1 + volatility / 2)
        lower = self.lower * (1 - volatility / 2)
        
        # 生成基础信号
        base_signals = np.select([
            (data['prediction'] > upper),
            (data['prediction'] < lower)
        ], [1, -1], default=0)
        
        # 过滤频繁信号（连续相同信号时归零）
        filtered = pd.Series(base_signals, index=data.index).diff()
        signals = filtered.where(filtered != 0, 0).astype(int)
        
        return signals